
import torch
import cv2
from one_layer_model import IlluNet_with_Quad
import os
import numpy as np
from basicsr.metrics import calculate_psnr, calculate_ssim
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
def calculate_mae(image1, image2):
    diff = np.abs(image1 - image2)
    mae = np.mean(diff)
    return mae
def to_deploy(mo):
    for m in mo.modules():
        if hasattr(m, 'switch_to_deploy'):
            m.switch_to_deploy()
    return mo


def run():
    model = IlluNet_with_Quad(3,3,None).eval()
    model.load_state_dict(torch.load(args.pre_weight))

    model = to_deploy(model).to(device)
    if not os.path.exists(args.out_dir):
        os.makedirs(args.out_dir)
    psnr_list = []
    ssim_list = []
    mse_list = []
    for filename in sorted( os.listdir(args.test_dir)):
        
        in_p = os.path.join(args.test_dir,filename)
        this_GT_p = os.path.join(args.GT_path,filename)
        img_gt = cv2.imread(this_GT_p).astype('float32')
        img_np = cv2.imread(in_p).astype('float32')/255.0
        img_ten = torch.from_numpy(img_np).permute(2,0,1).unsqueeze(0).to(device)
    
    
        with torch.no_grad():
            out = model(img_ten)
        enhanced_np = out.detach().cpu().squeeze(0).permute(1,2,0).numpy()*255.0
        psnr = calculate_psnr(enhanced_np,img_gt,crop_border=0)
        ssim = calculate_ssim(enhanced_np,img_gt,crop_border=0)
        mae = calculate_mae(enhanced_np,img_gt)
        
        psnr_list.append(psnr)
        ssim_list.append(ssim)
        mse_list.append(mae)
        print('{} , psnr {:.2f}  ssim {:.4f}  mae {:.4f}'.format(filename,psnr,ssim,mae))
        cv2.imwrite(args.out_dir+'/'+filename,enhanced_np)
    print('avg psnr {:.2f} ssim {:.4f} mae {:.2f} '.format(sum(psnr_list)/len(psnr_list),sum(ssim_list)/len(ssim_list),sum(mse_list)/len(mse_list)))

if __name__=='__main__':
    
    device = torch.device('cuda')   
    parser = ArgumentParser(description="validation script for refence",formatter_class=ArgumentDefaultsHelpFormatter)
    parser.add_argument('--pre_weight', default='./weight/sclm_ref.pth', type=str, help='weight of model')
    parser.add_argument('--GT_path', default='path to GT', type=str, help='path of ground truth')
    parser.add_argument('--out_dir', default="out/ref", type=str, help='output path')
    parser.add_argument('--test_dir', default="path to input", type=str, help='input path')

    args = parser.parse_args()
    run()
    
        



